{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6a8ca38c",
   "metadata": {},
   "source": [
    "## 定义测试数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "890ef889",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 1., 1., 1., 1., 1., 0., 0.],\n",
       "        [1., 1., 1., 1., 1., 1., 1., 0.]])"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第14章/虚拟数据\n",
    "import torch\n",
    "\n",
    "# 假设有2句话，8个词\n",
    "x = torch.ones(2, 8)\n",
    "\n",
    "# 两句话中各有一些pad\n",
    "x[0, 6:] = 0\n",
    "x[1, 7:] = 0\n",
    "\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "746a46e2",
   "metadata": {},
   "source": [
    "## PAD MASK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2f866dca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[False, False, False, False, False, False,  True,  True],\n",
       "        [False, False, False, False, False, False, False,  True]])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 第14章/定义key_padding_mask\n",
    "# key_padding_mask的定义方式，就是x中是pad的为True，否则是False\n",
    "key_padding_mask = x == 0\n",
    "\n",
    "key_padding_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "23ea0860",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[False, False, False, False, False, False, False, False],\n",
       "        [False, False, False, False, False, False, False, False],\n",
       "        [False, False, False, False, False, False, False, False],\n",
       "        [False, False, False, False, False, False, False, False],\n",
       "        [False, False, False, False, False, False, False, False],\n",
       "        [False, False, False, False, False, False, False, False],\n",
       "        [False, False, False, False, False, False, False, False],\n",
       "        [False, False, False, False, False, False, False, False]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第14章/定义encode_attn_mask\n",
    "#在encode阶段不需要定义encode_attn_mask\n",
    "#定义为None或者全False都可以\n",
    "encode_attn_mask = torch.ones(8, 8) == 0\n",
    "\n",
    "encode_attn_mask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65484574",
   "metadata": {},
   "source": [
    "## 上三角 MASK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ea205f63",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[False,  True,  True,  True,  True,  True,  True,  True],\n",
       "        [False, False,  True,  True,  True,  True,  True,  True],\n",
       "        [False, False, False,  True,  True,  True,  True,  True],\n",
       "        [False, False, False, False,  True,  True,  True,  True],\n",
       "        [False, False, False, False, False,  True,  True,  True],\n",
       "        [False, False, False, False, False, False,  True,  True],\n",
       "        [False, False, False, False, False, False, False,  True],\n",
       "        [False, False, False, False, False, False, False, False]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第14章/定义decode_attn_mask\n",
    "#在decode阶段需要定义decode_attn_mask\n",
    "#decode_attn_mask的定义方式是对角线以上为True的上三角矩阵\n",
    "decode_attn_mask = torch.tril(torch.ones(8, 8)) == 0\n",
    "\n",
    "decode_attn_mask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16cda128",
   "metadata": {},
   "source": [
    "## 词编码"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8961595a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],\n",
       " \n",
       "         [[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]]),\n",
       " torch.Size([2, 8, 12]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第14章/编码x\n",
    "x = x.unsqueeze(2)\n",
    "x = x.expand(-1, -1, 12)\n",
    "\n",
    "x, x.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "872f3445",
   "metadata": {},
   "source": [
    "## 多头注意力计算函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "54db78d4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([36, 12]),\n",
       " torch.Size([36]),\n",
       " torch.Size([12, 12]),\n",
       " torch.Size([12]))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第14章/定义multi_head_attention_forward()所需要的参数\n",
    "\n",
    "# in_proj就是qkv线性变换的参数\n",
    "in_proj_weight = torch.nn.Parameter(torch.randn(3 * 12, 12))\n",
    "in_proj_bias = torch.nn.Parameter(torch.zeros((3 * 12)))\n",
    "\n",
    "# out_proj就是输出时做线性变换的参数\n",
    "out_proj_weight = torch.nn.Parameter(torch.randn(12, 12))\n",
    "out_proj_bias = torch.nn.Parameter(torch.zeros(12))\n",
    "\n",
    "in_proj_weight.shape, in_proj_bias.shape, out_proj_weight.shape, out_proj_bias.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8541ee9c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([8, 2, 12]),\n",
       " tensor([[[0.1042, 0.2083, 0.2083, 0.2083, 0.1042, 0.2083, 0.0000, 0.0000],\n",
       "          [0.2083, 0.2083, 0.2083, 0.2083, 0.2083, 0.2083, 0.0000, 0.0000],\n",
       "          [0.2083, 0.2083, 0.1042, 0.2083, 0.2083, 0.1042, 0.0000, 0.0000],\n",
       "          [0.2083, 0.1042, 0.1042, 0.2083, 0.0000, 0.2083, 0.0000, 0.0000],\n",
       "          [0.1042, 0.2083, 0.1042, 0.2083, 0.1042, 0.2083, 0.0000, 0.0000],\n",
       "          [0.2083, 0.2083, 0.1042, 0.2083, 0.2083, 0.2083, 0.0000, 0.0000],\n",
       "          [0.2083, 0.2083, 0.2083, 0.2083, 0.2083, 0.2083, 0.0000, 0.0000],\n",
       "          [0.1042, 0.2083, 0.2083, 0.1042, 0.1042, 0.2083, 0.0000, 0.0000]],\n",
       " \n",
       "         [[0.1786, 0.0893, 0.1786, 0.0893, 0.0893, 0.0893, 0.1786, 0.0000],\n",
       "          [0.1786, 0.1786, 0.0893, 0.0893, 0.1786, 0.0893, 0.0893, 0.0000],\n",
       "          [0.1786, 0.1786, 0.1786, 0.1786, 0.1786, 0.1786, 0.0000, 0.0000],\n",
       "          [0.1786, 0.1786, 0.1786, 0.0893, 0.1786, 0.1786, 0.1786, 0.0000],\n",
       "          [0.0893, 0.0893, 0.1786, 0.1786, 0.0893, 0.1786, 0.1786, 0.0000],\n",
       "          [0.1786, 0.1786, 0.0893, 0.0893, 0.1786, 0.1786, 0.0893, 0.0000],\n",
       "          [0.1786, 0.1786, 0.0893, 0.1786, 0.1786, 0.0000, 0.0893, 0.0000],\n",
       "          [0.0893, 0.1786, 0.1786, 0.1786, 0.1786, 0.0893, 0.0893, 0.0000]]],\n",
       "        grad_fn=<MeanBackward1>),\n",
       " torch.Size([2, 8, 8]))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第14章/使用工具函数计算多头注意力\n",
    "data = {\n",
    "    #因为不是batch_first的，所以需要变形一下\n",
    "    'query': x.permute(1, 0, 2),\n",
    "    'key': x.permute(1, 0, 2),\n",
    "    'value': x.permute(1, 0, 2),\n",
    "    'embed_dim_to_check': 12,\n",
    "    'num_heads': 2,\n",
    "    'in_proj_weight': in_proj_weight,\n",
    "    'in_proj_bias': in_proj_bias,\n",
    "    'bias_k': None,\n",
    "    'bias_v': None,\n",
    "    'add_zero_attn': False,\n",
    "    'dropout_p': 0.2,\n",
    "    'out_proj_weight': out_proj_weight,\n",
    "    'out_proj_bias': out_proj_bias,\n",
    "    'key_padding_mask': key_padding_mask,\n",
    "    'attn_mask': encode_attn_mask,\n",
    "}\n",
    "\n",
    "score, attn = torch.nn.functional.multi_head_attention_forward(**data)\n",
    "\n",
    "score.shape, attn, attn.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "345a97bf",
   "metadata": {},
   "source": [
    "## 多头注意力层"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "24fa60b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([2, 8, 12]),\n",
       " tensor([[[0.1042, 0.2083, 0.2083, 0.2083, 0.2083, 0.1042, 0.0000, 0.0000],\n",
       "          [0.2083, 0.2083, 0.0000, 0.1042, 0.2083, 0.2083, 0.0000, 0.0000],\n",
       "          [0.2083, 0.2083, 0.1042, 0.1042, 0.2083, 0.1042, 0.0000, 0.0000],\n",
       "          [0.2083, 0.2083, 0.1042, 0.2083, 0.1042, 0.2083, 0.0000, 0.0000],\n",
       "          [0.2083, 0.2083, 0.1042, 0.1042, 0.2083, 0.2083, 0.0000, 0.0000],\n",
       "          [0.2083, 0.2083, 0.2083, 0.2083, 0.2083, 0.1042, 0.0000, 0.0000],\n",
       "          [0.0000, 0.2083, 0.2083, 0.1042, 0.2083, 0.1042, 0.0000, 0.0000],\n",
       "          [0.2083, 0.2083, 0.2083, 0.2083, 0.2083, 0.1042, 0.0000, 0.0000]],\n",
       " \n",
       "         [[0.1786, 0.1786, 0.1786, 0.1786, 0.1786, 0.1786, 0.1786, 0.0000],\n",
       "          [0.1786, 0.1786, 0.0893, 0.1786, 0.1786, 0.1786, 0.1786, 0.0000],\n",
       "          [0.0893, 0.0893, 0.0893, 0.0000, 0.1786, 0.1786, 0.1786, 0.0000],\n",
       "          [0.1786, 0.1786, 0.1786, 0.0893, 0.0893, 0.1786, 0.0893, 0.0000],\n",
       "          [0.1786, 0.0893, 0.1786, 0.1786, 0.1786, 0.1786, 0.0893, 0.0000],\n",
       "          [0.1786, 0.1786, 0.0893, 0.1786, 0.1786, 0.0893, 0.1786, 0.0000],\n",
       "          [0.0893, 0.0893, 0.1786, 0.1786, 0.0893, 0.1786, 0.0893, 0.0000],\n",
       "          [0.1786, 0.0893, 0.0893, 0.0893, 0.0893, 0.1786, 0.1786, 0.0000]]],\n",
       "        grad_fn=<MeanBackward1>),\n",
       " torch.Size([2, 8, 8]))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第14章/使用多头注意力工具层\n",
    "multihead_attention = torch.nn.MultiheadAttention(embed_dim=12,\n",
    "                                                    num_heads=2,\n",
    "                                                    dropout=0.2,\n",
    "                                                    batch_first=True)\n",
    "\n",
    "data = {\n",
    "    'query': x,\n",
    "    'key': x,\n",
    "    'value': x,\n",
    "    'key_padding_mask': key_padding_mask,\n",
    "    'attn_mask': encode_attn_mask,\n",
    "}\n",
    "\n",
    "score, attn = multihead_attention(**data)\n",
    "\n",
    "score.shape, attn, attn.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6202d9b",
   "metadata": {},
   "source": [
    "## 编码器层"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1801bd9f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 8, 12])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第14章/使用单层编码器工具层\n",
    "encoder_layer = torch.nn.TransformerEncoderLayer(\n",
    "    d_model=12,\n",
    "    nhead=2,\n",
    "    dim_feedforward=24,\n",
    "    dropout=0.2,\n",
    "    activation=torch.nn.functional.relu,\n",
    "    batch_first=True,\n",
    "    norm_first=True)\n",
    "\n",
    "data = {\n",
    "    'src': x,\n",
    "    'src_mask': encode_attn_mask,\n",
    "    'src_key_padding_mask': key_padding_mask,\n",
    "}\n",
    "\n",
    "out = encoder_layer(**data)\n",
    "\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b55fe8cf",
   "metadata": {},
   "source": [
    "## 编码器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "39a1c319",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 8, 12])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第14章/使用编码器工具层\n",
    "encoder = torch.nn.TransformerEncoder(\n",
    "    encoder_layer=encoder_layer,\n",
    "    num_layers=3,\n",
    "    norm=torch.nn.LayerNorm(normalized_shape=12))\n",
    "\n",
    "data = {\n",
    "    'src': x,\n",
    "    'mask': encode_attn_mask,\n",
    "    'src_key_padding_mask': key_padding_mask,\n",
    "}\n",
    "\n",
    "out = encoder(**data)\n",
    "\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eea48033",
   "metadata": {},
   "source": [
    "## 解码器层"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "eef4c26d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 8, 12])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第14章/使用单层解码器工具层\n",
    "decoder_layer = torch.nn.TransformerDecoderLayer(\n",
    "    d_model=12,\n",
    "    nhead=2,\n",
    "    dim_feedforward=24,\n",
    "    dropout=0.2,\n",
    "    activation=torch.nn.functional.relu,\n",
    "    batch_first=True,\n",
    "    norm_first=True)\n",
    "\n",
    "data = {\n",
    "    'tgt': x,\n",
    "    'memory': x,\n",
    "    'tgt_mask': decode_attn_mask,\n",
    "    'memory_mask': encode_attn_mask,\n",
    "    'tgt_key_padding_mask': key_padding_mask,\n",
    "    'memory_key_padding_mask': key_padding_mask,\n",
    "}\n",
    "\n",
    "out = decoder_layer(**data)\n",
    "\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8c8b314",
   "metadata": {},
   "source": [
    "## 解码器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "94f19f1f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 8, 12])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第14章/使用编码器工具层\n",
    "decoder = torch.nn.TransformerDecoder(\n",
    "    decoder_layer=decoder_layer,\n",
    "    num_layers=3,\n",
    "    norm=torch.nn.LayerNorm(normalized_shape=12))\n",
    "\n",
    "data = {\n",
    "    'tgt': x,\n",
    "    'memory': x,\n",
    "    'tgt_mask': decode_attn_mask,\n",
    "    'memory_mask': encode_attn_mask,\n",
    "    'tgt_key_padding_mask': key_padding_mask,\n",
    "    'memory_key_padding_mask': key_padding_mask,\n",
    "}\n",
    "\n",
    "out = decoder(**data)\n",
    "\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12e203f7",
   "metadata": {},
   "source": [
    "## Transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "342fb1b5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 8, 12])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#第14章/使用Transformer工具模型\n",
    "transformer = torch.nn.Transformer(d_model=12,\n",
    "                                    nhead=2,\n",
    "                                    num_encoder_layers=3,\n",
    "                                    num_decoder_layers=3,\n",
    "                                    dim_feedforward=24,\n",
    "                                    dropout=0.2,\n",
    "                                    activation=torch.nn.functional.relu,\n",
    "                                    custom_encoder=encoder,\n",
    "                                    custom_decoder=decoder,\n",
    "                                    batch_first=True,\n",
    "                                    norm_first=True)\n",
    "\n",
    "data = {\n",
    "    'src': x,\n",
    "    'tgt': x,\n",
    "    'src_mask': encode_attn_mask,\n",
    "    'tgt_mask': decode_attn_mask,\n",
    "    'memory_mask': encode_attn_mask,\n",
    "    'src_key_padding_mask': key_padding_mask,\n",
    "    'tgt_key_padding_mask': key_padding_mask,\n",
    "    'memory_key_padding_mask': key_padding_mask,\n",
    "}\n",
    "\n",
    "out = transformer(**data)\n",
    "\n",
    "out.shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
